import datetime
from sqlalchemy import text
import pandas as pd
from retrying import retry
import tushare as ts
# import sys
# sys.path.append('E:\\gitee\\mp_alwaysup_fastapi')
from process import db_connect as dbConn


@retry(wait_fixed=1000 * 10,  # 每次重试等待20分钟
       stop_max_delay=5 * 60 * 1000,  # 最大重试时间为5分钟
       retry_on_exception=lambda e: True,  # 任何异常都触发重试
       retry_on_result=lambda result: result is None  # 返回None时触发重试
       )
def get_tushare_pro():
    # 设置tushare.token
    ts.set_token("38bb3cd1b6af2d75a7d7e506db8fd60354168642b400fa2104af81c5")
    try:
        pro = ts.pro_api()
        return pro
    except Exception as e:
        print(f"Error connect tushare pro: {e}")


class tushare:
    def __init__(self):
        # 创建数据库引擎
        self.pro = get_tushare_pro()
        self.engine = None
        # 获取当前日期
        self.today = datetime.datetime.now().strftime("%Y%m%d")

    @retry(wait_fixed=5000, stop_max_attempt_number=10)
    def connect(self):
        try:
            if not self.engine:
                self.engine = dbConn.get_Mpdb_engine()
        except Exception as e:
            print(f"Error connect engine: {e}")

    # 股票列表
    @retry(wait_fixed=5000, stop_max_attempt_number=10)
    def get_stock_basic_info(self):
        try:
            data1 = self.pro.stock_basic(
                exchange="", list_status='D',
                fields="ts_code,symbol,name,area,industry,fullname,enname,cnspell,market,exchange,curr_type,list_status,list_date,delist_date,is_hs,act_name,act_ent_type",
            )
            data2 = self.pro.stock_basic(
                exchange="", list_status='L',
                fields="ts_code,symbol,name,area,industry,fullname,enname,cnspell,market,exchange,curr_type,list_status,list_date,delist_date,is_hs,act_name,act_ent_type",
            )
            stockBasicDf = pd.concat([data1, data2], ignore_index=True)
            print("股票基本信息获取成功")
            return stockBasicDf

        except Exception as e:
            print(f"获取股票基本信息失败: {e}")

    # 判断是否交易日

    @retry(wait_fixed=1000 * 10,  # 每次重试等待10秒钟
           stop_max_delay=5 * 60 * 1000,  # 最大重试时间为5分钟
           retry_on_exception=lambda e: True,  # 任何异常都触发重试
           retry_on_result=lambda result: result is None  # 返回None时触发重试
           )
    def is_trading_day(self, date):
        try:
            cal = self.pro.trade_cal(start_date=date, end_date=date)
            if cal.empty:
                return False
            return cal.iloc[0]["is_open"] == 1
        except Exception as e:
            print(f"Exception: {e}")

    @retry(wait_fixed=5000, stop_max_attempt_number=10)
    def get_trade_cal(self):
        try:
            tradeDate = datetime.datetime.now().strftime("%Y%m%d")
            df = self.pro.trade_cal(
                exchange="", start_date='', end_date=tradeDate)
            return df
        except Exception as e:
            print(f"Error fetching trade calendar: {e}")

    @retry(wait_fixed=1000*10, stop_max_attempt_number=10)
    def get_daily_info(self, tradeDate):
        try:
            df = self.pro.daily_info(exchange="SZ,SH", trade_date=tradeDate)
            return df
        except Exception as e:
            print(f"Error fetching daily_info: {e}")

    @retry(wait_fixed=1000*10, stop_max_attempt_number=10)
    def get_daily_basic(self, tradeDate):
        try:
            df = self.pro.daily_basic(trade_date=tradeDate)
            return df
        except Exception as e:
            print(f"Error fetching daily_basic: {e}")

    @retry(wait_fixed=1000*10, stop_max_attempt_number=10)
    def get_daily_basic_mv_count(self, tradeDate):
        if not self.engine:
            self.connect()
        try:
            sql = text(f"SELECT "
                       f"trade_date,  COUNT(*) AS total_stocks,"
                       f"SUM(IF(circ_mv >= 0 AND circ_mv < 50000, 1, 0)) AS mv_5y, "
                       f"SUM(IF(circ_mv >= 50000 AND circ_mv < 300000, 1, 0)) AS mv_5y_30y, "
                       f"SUM(IF(circ_mv >= 300000 AND circ_mv < 1000000, 1, 0)) AS mv_30y_1by, "
                       f"SUM(IF(circ_mv >= 1000000 AND circ_mv < 5000000, 1, 0)) AS mv_1by_5by, "
                       f"SUM(IF(circ_mv >= 5000000 AND circ_mv < 10000000, 1, 0)) AS mv_5by_1ky, "
                       f"SUM(IF(circ_mv >= 10000000 AND circ_mv < 50000000, 1, 0)) AS mv_1ky_5ky, "
                       f"SUM(IF(circ_mv >= 50000000 AND circ_mv < 100000000, 1, 0)) AS mv_5ky_1wy, "
                       f"SUM(IF(circ_mv >= 100000000 AND circ_mv < 200000000, 1, 0)) AS mv_1wy_2wy,"
                       f"SUM(IF(circ_mv >= 200000000, 1, 0))  AS mv_above2wy "
                       f"FROM daily_basic "
                       f"WHERE trade_date = '{tradeDate}' "
                       f"GROUP BY trade_date")
            result = pd.read_sql(sql, self.engine)
        except Exception as e:
            print("Database error:", e)
            result = None
        return result

    @retry(wait_fixed=1000*60, stop_max_attempt_number=10)
    def get_kday_daily_all(self, tradeDate):
        try:
            df = self.pro.daily(trade_date=tradeDate)
            return df
        except Exception as e:
            print(f"Error fetching daily k-line data: {e}")

# -------------------------------------------------------------
    @retry(wait_fixed=1000 * 10,  # 每次重试等待20分钟
           stop_max_delay=5 * 60 * 1000,  # 最大重试时间为5分钟
           retry_on_exception=lambda e: True,  # 任何异常都触发重试
           retry_on_result=lambda result: result is None  # 返回None时触发重试
           )
    def get_trade_date_list(self):
        if not self.engine:
            self.connect()

        trade_date_list = []
        cal_data = None

        try:
            sql = text("SELECT * FROM trade_cal WHERE is_open = 1 order by  cal_date  desc")
            cal_data = pd.read_sql(sql, self.engine)
            if not cal_data.empty:
                trade_date_list = cal_data["cal_date"].tolist()
                # print('from db')
            else:
                print("Empty result from trade_cal table")
        except Exception as e:
            print("Database error:", e)
            try:
                # 使用Tushare作为备选方案获取交易日历数据
                cal_data = self.pro.trade_cal(
                    exchange="", is_open="1", end_date=self.today)
                if not cal_data.empty:
                    trade_date_list = cal_data["cal_date"].tolist()
                    # print('from ts')
                else:
                    print("Empty result from Tushare API")
            except Exception as e:
                print("Tushare get data error:", e)
                return None, None

        trade_date_list.sort()
        return trade_date_list, cal_data

    @retry(wait_fixed=5000, stop_max_attempt_number=10)
    def get_missing_trade_days(self, table_name):
        try:
            if not self.engine:
                self.connect()
            # 查询已有的交易日期
            sql = f"SELECT DISTINCT trade_date FROM {table_name} order by trade_date"
            df1 = pd.read_sql(sql, self.engine)

            # 获取所有可能的交易日期列表
            trade_date_list, cal_data = self.get_trade_date_list()
            df2 = pd.DataFrame(trade_date_list,
                               columns=['trade_date'])

            # 统一日期格式
            df1['trade_date'] = pd.to_datetime(
                df1['trade_date']).dt.strftime('%Y%m%d')
            df2['trade_date'] = pd.to_datetime(
                df2['trade_date']).dt.strftime('%Y%m%d')
            # 合并并去除重复项
            df_merge = pd.concat([df1, df2], ignore_index=True)
            result = df_merge.drop_duplicates(subset=['trade_date'], keep=False)[
                'trade_date'].tolist()
        except Exception as e:
            # 异常处理，记录日志或返回空列表
            print(f"Error occurred while fetching {table_name} trade days:", e)
            result = []
        return result

    def unGet_tushare_data_tradeDay_list(self, dbName):
        tradeDay_list = self.get_missing_trade_days(dbName)
        if dbName == 'market_amount':
            tradeDay_list = [
                day for day in tradeDay_list if int(day) > 20071128]
        return tradeDay_list


tushare = tushare()
# print(tushare.get_trade_date_list())
# trade_date_list, cal_date = tushare.get_trade_date_list()
# from datetime import date, datetime, timedelta
# previous_date = datetime.now().date() - timedelta(days=4)
# print(previous_date)
# print(tushare.is_trading_day(previous_date))
# print(tushare.is_trading_day('20240811'))
# tushare.unGet_market_amount_tradeDay_list()
